# -*- coding: utf-8 -*-
import os
import unittest

import h5py
import hdf5plugin  # noqa: F401
import numpy as np
from astropy.io import fits


class Fits2Hdf5Test(unittest.TestCase):
    def setUp(self):
        from f2h5.config import DATA_DIR
        filename = "FRB20240114A_tracking-M01_0001"
        self.fits_hdus = fits.open(os.path.join(DATA_DIR, f"{filename}.fits"))
        self.h5_root = h5py.File(os.path.join(DATA_DIR, f"{filename}.h5"), "r")[filename]

    def test_two_formats_reads_identical(self):
        for hdu in self.fits_hdus:
            hdu_name = hdu.name
            h5_grp = self.h5_root[hdu_name]
            h5_attrs = {k: v for k, v in h5_grp.attrs.items() if not (k.startswith("__") and k.startswith("__"))}
            fits_attrs = dict(hdu.header)
            if (comment := fits_attrs.pop("COMMENT", None)) is not None:
                fits_attrs["COMMENT"] = str(comment)
            self.assertEqual(h5_attrs, fits_attrs)

            for ds_name in h5_grp:
                fits_ds = hdu.data.field(ds_name)
                h5_ds = h5_grp[ds_name][...]
                if np.issubdtype(h5_ds.dtype, np.floating):
                    self.assertTrue(np.allclose(fits_ds, h5_ds))
                else:
                    self.assertTrue(np.all(np.equal(fits_ds.ravel(), h5_ds.ravel())))


if __name__ == "__main__":
    unittest.main()
